package com.youxiake.lock.zk;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

public class ZooKeeperDistributedLock implements Watcher {

    private final ZooKeeper zk;
    private final String locksRoot = "/locks";
    private final String productId;
    private String waitNode;
    private String lockNode;
    private CountDownLatch latch;
    private final CountDownLatch connectedLatch = new CountDownLatch(1);
    private final int sessionTimeout = 300000;

    public static void main(String[] args) throws InterruptedException {
        ZooKeeperDistributedLock lock = new ZooKeeperDistributedLock("1");
        for (int i = 0; i < 4; i++) {
            new Thread(() -> {
                String key = lock.acquireDistributedLock();
                System.out.println("加锁结果" + key);
                try {
                    Thread.sleep(3000);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
                lock.unlock(key);
            }).start();
        }
    }

    public ZooKeeperDistributedLock(String productId) {
        this.productId = productId;
        try {
            String address = "1.116.186.202:2181";
            zk = new ZooKeeper(address, sessionTimeout, this);
            connectedLatch.await();
        } catch (IOException | InterruptedException e) {
            throw new LockException(e);
        }
    }

    @Override
    public void process(WatchedEvent event) {
        if (event.getState() == Event.KeeperState.SyncConnected) {
            connectedLatch.countDown();
            return;
        }

        if (this.latch != null) {
            System.out.println("释放了锁");
            this.latch.countDown();
        }
    }

    public String acquireDistributedLock() {
        try {
            String key = this.tryLock();
            System.out.printf("key %s, lockNode %s%n", key, lockNode);
            if (key != null && !"".equals(key) && key.equals("/locks/" + lockNode)) {
                System.out.printf("success to acquire lock , key: %s%n", key);
                return key;
            } else {
                waitForLock(key, waitNode, sessionTimeout);
                System.out.printf("success to acquire lock , key: %s%n", key);
            }
            return key;
        } catch (KeeperException | InterruptedException e) {
            throw new LockException(e);
        }
    }

    public String tryLock() {
        try {
            // 传入进去的locksRoot + “/” + productId
            // 假设productId代表了一个商品id，比如说1
            // locksRoot = locks
            // /locks/10000000000，/locks/10000000001，/locks/10000000002
            String node = zk.create(locksRoot + "/" + productId, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);

            // 看看刚创建的节点是不是最小的节点
            // locks：10000000000，10000000001，10000000002
            List<String> locks = zk.getChildren(locksRoot, false);
            Collections.sort(locks);
            lockNode = locks.get(0);
            if(node.equals(locksRoot+"/"+ locks.get(0))){
                //如果是最小的节点,则表示取得锁
                return node;
            }

            //如果不是最小的节点，找到比自己小1的节点
            int previousLockIndex = -1;
            for(int i = 0; i < locks.size(); i++) {
                if(node.equals(locksRoot + "/" + locks.get(i))) {
                    previousLockIndex = i - 1;
                    break;
                }
            }

            this.waitNode = locks.get(0);

            return node;
        } catch (KeeperException | InterruptedException e) {
            throw new LockException(e);
        }
    }

    private boolean waitForLock(String key, String waitNode, long waitTime) throws InterruptedException, KeeperException {
        Stat stat = zk.exists(locksRoot + "/" + waitNode, true);
        if (stat != null) {
            if (this.latch == null) {
                this.latch = new CountDownLatch(1);
            }
            this.latch.await(waitTime, TimeUnit.MILLISECONDS);

            List<String> locks = zk.getChildren(locksRoot, false);
            Collections.sort(locks);
            if(key.equals(locksRoot+"/"+ locks.get(0))){
                //如果是最小的节点,则表示取得锁
                return true;
            } else {
                waitForLock(key, locks.get(0), waitTime);
            }
        }
        return true;
    }

    public void unlock(String key) {
        try {
            // 删除/locks/10000000000节点
            // 删除/locks/10000000001节点
            System.out.println("unlock " + key);
            zk.delete(key, -1);
            this.latch.countDown();
        } catch (InterruptedException | KeeperException e) {
            e.printStackTrace();
        }
    }

    public static class LockException extends RuntimeException {
        private static final long serialVersionUID = 1L;

        public LockException(String e) {
            super(e);
        }

        public LockException(Exception e) {
            super(e);
        }
    }
}